import os
import json
from typing import Annotated

from langchain_core.messages import BaseMessage
from typing_extensions import TypedDict

from langgraph.graph import StateGraph, START
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool


@tool
def calculate(expression: str) -> str:
    """🧮 数学计算工具
    
    功能：执行数学计算并返回结果
    
    适用场景：
    - 基础数学运算：加减乘除（+、-、*、/）
    - 支持括号优先级计算
    - 支持小数运算
    - 支持复杂数学表达式
    
    参数：
    - expression (str): 需要计算的数学表达式，如 "25 * 4 + 10" 或 "(100 + 50) / 3"
    
    返回：
    - str: 计算结果，格式为 "计算结果：表达式 = 结果"
    
    使用示例：
    - calculate("25 * 4 + 10") -> "计算结果：25 * 4 + 10 = 110"
    - calculate("100 / 5") -> "计算结果：100 / 5 = 20.0"
    """
    print("calculate 工具被调用, 进行计算")
    try:
        # 安全的数学计算
        allowed_chars = set('0123456789+-*/().,e ')
        if not all(c in allowed_chars for c in expression):
            return "错误：表达式包含不允许的字符"
        
        result = eval(expression)
        return f"计算结果：{expression} = {result}"
    except Exception as e:
        return f"计算错误：{str(e)}"


@tool
def get_weather(location: str) -> str:
    """🌤️ 天气查询工具
    
    功能：获取指定城市的实时天气信息
    
    适用场景：
    - 查询任何地点的天气状况
    - 获取温度、天气描述、空气质量等信息
    - 支持中文城市名称（如：北京、上海、广州、深圳）
    - 支持英文城市名称（如：Beijing、Shanghai、New York）
    
    参数：
    - location (str): 城市名称，支持中英文，如 "北京"、"上海"、"Beijing"
    
    返回：
    - str: 天气信息，包含城市、天气状况、温度、空气质量等
    
    使用示例：
    - get_weather("北京") -> "北京 今天天气晴天☀️, 温度26°C, 空气质量良好"
    - get_weather("上海") -> "上海 今天天气晴天☀️, 温度26°C, 空气质量良好"
    
    触发关键词：天气、温度、气候、下雨、晴天、阴天、空气质量等
    """
    print("get_weather 工具被调用, 进行天气查询")

    info = f"{location} 今天天气晴天☀️, 温度26°C, 空气质量良好"
    return info

llm = ChatOpenAI(
    model_name="deepseek-chat",
    openai_api_key="sk-c2acb26542994445a95052fbcae9cd02",
    openai_api_base="https://api.deepseek.com/v1",
)



class State(TypedDict):
    messages: Annotated[list, add_messages]

# --- Graph Definition ---
graph_builder = StateGraph(State)


tools = [calculate, get_weather]

llm_with_tools = llm.bind_tools(tools)

# Define the chatbot node
def chatbot(state: State):
    return {"messages": [llm_with_tools.invoke(state["messages"])]}

# Add nodes to the graph
graph_builder.add_node("chatbot", chatbot)
tool_node = ToolNode(tools=tools)
graph_builder.add_node("tools", tool_node)

# Define the graph's flow
graph_builder.add_edge(START, "chatbot")
graph_builder.add_conditional_edges(
    "chatbot",
    tools_condition,
)
graph_builder.add_edge("tools", "chatbot")

# Compile the graph
graph = graph_builder.compile()


# --- Run the Chatbot ---
# This is the missing part that starts the conversation loop.
def run_chatbot():
    print("Chatbot started! Ask a question that requires a search, like 'What is LangGraph?'.")
    print("Type 'quit', 'exit', or 'q' to stop.")
    while True:
        try:
            user_input = input("User: ")
            if user_input.lower() in ["quit", "exit", "q"]:
                print("Goodbye!")
                break
            
            # Stream the graph execution
            for event in graph.stream({"messages": [("user", user_input)]}):
                for node_name, value in event.items():
                    print(f"---Output from {node_name}---")
                    # Only print the final assistant message to avoid clutter
                    if node_name == "chatbot" and not value['messages'][-1].tool_calls:
                         print("Assistant:", value["messages"][-1].content)
                    elif node_name == "tools":
                        print(f"Tool Output: {value['messages'][0].content}")


        except (KeyboardInterrupt, EOFError):
            print("\nGoodbye!")
            break
        except Exception as e:
            import traceback
            print(f"An error occurred: {e}")
            traceback.print_exc()

# This line actually executes the chatbot loop when you run the script.
if __name__ == "__main__":
    run_chatbot()
